/*
 * TreeFactorizer.java
 *
 * Created on July 13, 2006, 1:51 PM
 *
 */

package galronnlp.util;

import galronnlp.perceptron.Node;
import java.io.BufferedReader;
import java.io.FileReader;
import java.util.Iterator;
import java.util.LinkedList;

import galronnlp.io.PennTreeReader;

/**
 * The <code>TreeFactorizer</code> class contains methods to factorize n-ary trees.
 *
 *
 * @author Daniel A. Galron
 */
public class TreeFactorizer {
    
    public static int TREE = 0;
    public static int NODE = 1;
    
    SymbolGenerator gen = new SymbolGenerator();
    
    
    /** Creates a new instance of TreeFactorizer */
    public TreeFactorizer() {
    }
    
    /** 
     * Creates a new instance of TreeFactorizer 
     * @params gen the <code>SymbolGenerator</code> to use in factoring trees.
     */
    public TreeFactorizer(SymbolGenerator gen) {
        this.gen = gen;
    }
    
/*    public Node rightFactorizeCNF(Node n) {
        if(n.type == Node.TERMINAL) {
            return n;
        } else if(n.type == Node.POS) {
            return n;
        } else {
            if(n.children.size() == 2) {
                Node left = rightFactorizeCNF(n.children.getFirst());
                Node right = rightFactorizeCNF(n.children.getLast());
                Node ret = new Node(n);
                ret.addChild(left);
                ret.addChild(right);
                return ret;
            } else if(n.children.size() == 1) {
                Node child = rightFactorizeCNF(n.children.getFirst());
                String symbolString = child.value().toString();
                Symbol tempSymbol = gen.create(n.value() + ":" + symbolString);
                Node ret = new Node(n);
                ret.addChild(child);
                ret.value = tempSymbol;
                return ret;
            } else {
                Symbol parent = n.value();
                LinkedList<Node> childNodes = new LinkedList<Node>();
                for(Iterator<Node> i = n.children.iterator(); i.hasNext();) {
                    childNodes.addLast(rightFactorizeCNF(i.next()));
                }
                Node newNode = childNodes.getLast();
                String symbolString = newNode.value().toString();
                for(int i = childNodes.size() - 2; i > 0; i--) {
                    Node nextNode = childNodes.get(i);
                    symbolString = nextNode.value().toString() + "+" + symbolString;
                    Symbol tempSymbol = gen.create(parent.toString() + ":" + symbolString);
                    
                    newTree = new BinaryTree(tempSymbol, nextTree, newTree, probability);
                }
                newTree = new BinaryTree(parent, children.get(0), newTree, tree.prob());                
            }
        }
    }
    */
    /**
     * Factors an n-ary tree into a right-branching CNF tree.
     * @params tree the n-ary tree to factor
     * @return a new, right-branching CNF tree
     */
    public Tree rightFactorizeCNF(Tree tree) {
        if(tree.type() == Tree.LEXICAL) {
            return tree;
        } else if(tree.type() == Tree.BINARY) {
            Tree newTree = new BinaryTree(tree.cat(), this.rightFactorizeCNF((Tree)tree.children.getFirst()), this.rightFactorizeCNF((Tree)tree.children.getLast()), tree.prob());
            return newTree;
        } else if(tree.type() == Tree.UNARY) {
            Tree newTree = this.rightFactorizeCNF((Tree)tree.children.getFirst());
            String symbolString = newTree.cat().toString();
            Symbol tempSymbol = gen.create(tree.cat() + ":" + symbolString);
            newTree.category = tempSymbol;
            return newTree;
        } else {
            Symbol parent = tree.cat();
            LinkedList<Tree> children = new LinkedList<Tree>();
            for(Iterator<Tree> i = tree.children.iterator(); i.hasNext(); ) {
                children.addLast(rightFactorizeCNF(i.next()));
            }
            Tree newTree = children.getLast();
            String symbolString = newTree.cat().toString();
            for(int i = children.size() - 2; i > 0; i--) {
                Tree nextTree = children.get(i);
                symbolString = nextTree.cat().toString() + "+" + symbolString;
                Symbol tempSymbol = gen.create(parent.toString() + ":" + symbolString);
                double probability = nextTree.prob() * newTree.prob();
                newTree = new BinaryTree(tempSymbol, nextTree, newTree, probability);
            }
            newTree = new BinaryTree(parent, children.get(0), newTree, tree.prob());
            return newTree;
        }
    }

    
    /**
     * Factors an n-ary tree into a left-branching CNF tree.
     * @params tree the n-ary tree to factor
     * @return a new, right-branching CNF tree
     */    
    public Tree leftFactorizeCNF(Tree tree) {
        if(tree.type() == Tree.LEXICAL) {
            return tree;
        } else if(tree.type() == Tree.UNARY) {
            Tree newTree = this.rightFactorizeCNF((Tree)tree.children.getFirst());
            newTree.category = gen.create(tree.cat() + ":" + newTree.cat());
            return newTree;
        } else if(tree.type() == Tree.BINARY) {
            Tree newTree = new BinaryTree(tree.cat(), this.leftFactorizeCNF((Tree)tree.children.getFirst()), this.rightFactorizeCNF((Tree)tree.children.getLast()), tree.prob());
            return newTree;
        } else {
            Symbol parent = tree.cat();
            LinkedList<Tree> children = new LinkedList<Tree>();
            for(Iterator<Tree> i = tree.children.iterator(); i.hasNext(); ) {
                children.addLast(leftFactorizeCNF(i.next()));
            }
            Tree newTree = children.getFirst();
            String symbolString = newTree.cat().toString();
            for(int i = 1; i < children.size() - 1; i++) {
                Tree nextTree = children.get(i);
                symbolString = symbolString + "+" + nextTree.cat().toString();
                Symbol newSymbol = gen.create(parent.toString() + ":" + symbolString);
                double probability = newTree.prob() * nextTree.prob();
                newTree = new BinaryTree(newSymbol, newTree, nextTree, probability);
            }
            newTree = new BinaryTree(parent, newTree, children.get(children.size() - 1), tree.prob());
            return newTree;
        }
    }
    
    public Tree rightMarkovFactorizeCNF(Tree tree, int order) {
        if(tree.type() == Tree.LEXICAL) {
            return tree;
        } else if(tree.type() == Tree.BINARY) {
            Tree newTree = new BinaryTree(tree.cat(), this.rightMarkovFactorizeCNF((Tree)tree.children.getFirst(), order), this.rightMarkovFactorizeCNF((Tree)tree.children.getLast(), order), tree.prob());
            return newTree;
        } else if(tree.type() == Tree.UNARY) {
            Tree newTree = this.rightMarkovFactorizeCNF((Tree)tree.children.getFirst(), order);
            
            return newTree;
        } else {
            Symbol parent = tree.cat();
            LinkedList<Tree> children = new LinkedList<Tree>();
            for(Iterator<Tree> i = tree.children.iterator(); i.hasNext(); ) {
                children.addLast(rightMarkovFactorizeCNF(i.next(), order));
            }
            
            Tree newTree = children.getLast();
            for(int i = children.size() - 2; i > 0; i--) {
                Tree nextTree = children.get(i);
                String newSymbolString = nextTree.cat().toString();
                String[] tempSymTok = newTree.cat().toString().split(":");
                if(tempSymTok.length == 2) {
                    String[] symbols = tempSymTok[1].split("\\+");
                    if(order > symbols.length) {
                        newSymbolString = newSymbolString + "+" + tempSymTok[1];
                    } else {
                        for(int j = 0; j < order - 1; j++) {
                            newSymbolString = newSymbolString + "+" + symbols[j];
                        }
                    }
                }
                else {
                    if(order > 1) {
                        newSymbolString = newSymbolString + "+" + tempSymTok[0];
                    }
                }
                Symbol newSymbol = gen.create(parent.toString() + ":" + newSymbolString);
                double probability = nextTree.prob() * newTree.prob();
                newTree = new BinaryTree(newSymbol, nextTree, newTree, probability);                
            }
            newTree = new BinaryTree(parent, children.get(0), newTree, tree.prob());
            return newTree;
        }        
    }
    
    public static void main(String[] args) {
        String treeFile = args[0];
        SymbolGenerator gen = new SymbolGenerator();
        TreeFactorizer tf = new TreeFactorizer(gen);
        try {
            PennTreeReader reader = new PennTreeReader(new BufferedReader(new FileReader(treeFile)), gen);
            Tree tree = reader.readTree();
            System.out.println(tree);
            System.out.println(tf.rightFactorizeCNF(tree));
        } catch(Exception e) {
            e.printStackTrace();
        }
        
    }
}
